#include <iostream>
#include <opencv2/opencv.hpp>
#include <fstream>
#include <iomanip>
#include <filesystem>
#include <vector>
#include <string>
#include <cmath>

#include "co_detr_ros/tensorrt_detector.h"

namespace fs = std::filesystem;

// 计算两个向量的差异统计
void analyzeDifferences(const std::vector<float>& fp16_results, const std::vector<float>& fp32_results) {
    if (fp16_results.size() != fp32_results.size()) {
        std::cout << "❌ Result sizes don't match: FP16=" << fp16_results.size() 
                  << ", FP32=" << fp32_results.size() << std::endl;
        return;
    }
    
    std::vector<float> differences;
    float max_diff = 0.0f;
    float min_diff = 0.0f;
    float sum_diff = 0.0f;
    float sum_squared_diff = 0.0f;
    
    for (size_t i = 0; i < fp16_results.size(); ++i) {
        float diff = fp32_results[i] - fp16_results[i];
        differences.push_back(diff);
        
        if (diff > max_diff) max_diff = diff;
        if (diff < min_diff) min_diff = diff;
        sum_diff += diff;
        sum_squared_diff += diff * diff;
    }
    
    float mean_diff = sum_diff / fp16_results.size();
    float variance = (sum_squared_diff / fp16_results.size()) - (mean_diff * mean_diff);
    float std_dev = std::sqrt(variance);
    
    std::cout << "\n=== Precision Analysis ===" << std::endl;
    std::cout << "Total elements: " << fp16_results.size() << std::endl;
    std::cout << "Max difference: " << max_diff << std::endl;
    std::cout << "Min difference: " << min_diff << std::endl;
    std::cout << "Mean difference: " << mean_diff << std::endl;
    std::cout << "Standard deviation: " << std_dev << std::endl;
    std::cout << "RMS difference: " << std::sqrt(sum_squared_diff / fp16_results.size()) << std::endl;
    
    // 统计差异分布
    int small_diff_count = 0;
    int medium_diff_count = 0;
    int large_diff_count = 0;
    
    for (float diff : differences) {
        float abs_diff = std::abs(diff);
        if (abs_diff < 0.001f) small_diff_count++;
        else if (abs_diff < 0.01f) medium_diff_count++;
        else large_diff_count++;
    }
    
    std::cout << "\n=== Difference Distribution ===" << std::endl;
    std::cout << "Small differences (<0.001): " << small_diff_count 
              << " (" << (small_diff_count * 100.0f / fp16_results.size()) << "%)" << std::endl;
    std::cout << "Medium differences (0.001-0.01): " << medium_diff_count 
              << " (" << (medium_diff_count * 100.0f / fp16_results.size()) << "%)" << std::endl;
    std::cout << "Large differences (>0.01): " << large_diff_count 
              << " (" << (large_diff_count * 100.0f / fp16_results.size()) << "%)" << std::endl;
}

// 比较检测结果
void compareDetections(const std::vector<float>& fp16_results, const std::vector<float>& fp32_results) {
    std::cout << "\n=== Detection Comparison ===" << std::endl;
    
    // 应用NMS到两个结果
    TensorRTDetector temp_detector("", "");
    std::vector<float> fp16_nms = temp_detector.applyNMS(fp16_results, 0.3f, 0.5f);
    std::vector<float> fp32_nms = temp_detector.applyNMS(fp32_results, 0.3f, 0.5f);
    
    std::cout << "FP16 detections after NMS: " << fp16_nms.size() / 5 << std::endl;
    std::cout << "FP32 detections after NMS: " << fp32_nms.size() / 5 << std::endl;
    
    // 比较检测框
    size_t min_detections = std::min(fp16_nms.size() / 5, fp32_nms.size() / 5);
    int matching_boxes = 0;
    int confidence_diff_count = 0;
    
    for (size_t i = 0; i < min_detections; ++i) {
        size_t fp16_idx = i * 5;
        size_t fp32_idx = i * 5;
        
        // 比较坐标（允许小的误差）
        bool coords_match = true;
        for (int j = 0; j < 4; ++j) {
            if (std::abs(fp16_nms[fp16_idx + j] - fp32_nms[fp32_idx + j]) > 0.01f) {
                coords_match = false;
                break;
            }
        }
        
        if (coords_match) {
            matching_boxes++;
        }
        
        // 检查置信度差异
        float conf_diff = std::abs(fp16_nms[fp16_idx + 4] - fp32_nms[fp32_idx + 4]);
        if (conf_diff > 0.05f) {
            confidence_diff_count++;
        }
    }
    
    std::cout << "Matching bounding boxes: " << matching_boxes << "/" << min_detections 
              << " (" << (matching_boxes * 100.0f / min_detections) << "%)" << std::endl;
    std::cout << "Detections with confidence difference >0.05: " << confidence_diff_count 
              << "/" << min_detections << std::endl;
}

int main(int argc, char** argv) {
    // Configuration paths
    std::string fp16_model_path = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/models/co_detr_layernorm_split_fp16.plan";
    std::string fp32_model_path = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/models/co_detr_layernorm_split_fp32.plan";
    std::string plugin_path = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/grid_sampler/build/libtrtgrid_sampler.so";
    std::string test_image = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/data/scene-1_000003.jpg";
    
    std::cout << "=== Co-DETR FP16 vs FP32 Engine Comparison ===" << std::endl;
    
    // Test FP16 engine
    std::cout << "\n--- Testing FP16 Engine ---" << std::endl;
    TensorRTDetector fp16_detector(fp16_model_path, plugin_path);
    if (!fp16_detector.initialize()) {
        std::cerr << "Failed to initialize FP16 detector" << std::endl;
        return -1;
    }
    
    std::vector<float> fp16_results = fp16_detector.inferenceFromImage(test_image);
    if (fp16_results.empty()) {
        std::cerr << "FP16 inference failed" << std::endl;
        return -1;
    }
    
    // Test FP32 engine
    std::cout << "\n--- Testing FP32 Engine ---" << std::endl;
    TensorRTDetector fp32_detector(fp32_model_path, plugin_path);
    if (!fp32_detector.initialize()) {
        std::cerr << "Failed to initialize FP32 detector" << std::endl;
        return -1;
    }
    
    std::vector<float> fp32_results = fp32_detector.inferenceFromImage(test_image);
    if (fp32_results.empty()) {
        std::cerr << "FP32 inference failed" << std::endl;
        return -1;
    }
    
    // Analyze differences
    analyzeDifferences(fp16_results, fp32_results);
    compareDetections(fp16_results, fp32_results);
    
    // Save individual visualizations
    std::string fp16_viz_path = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/results/fp16_detection_result.jpg";
    std::string fp32_viz_path = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/results/fp32_detection_result.jpg";
    
    fp16_detector.visualizeResults(fp16_results, fp16_viz_path);
    fp32_detector.visualizeResults(fp32_results, fp32_viz_path);
    
    std::cout << "\n=== Comparison Complete ===" << std::endl;
    std::cout << "FP16 visualization: " << fp16_viz_path << std::endl;
    std::cout << "FP32 visualization: " << fp32_viz_path << std::endl;
    
    return 0;
} 